
#include "common.h"

#include <linux/bpf.h>
#include <bpf/bpf_endian.h>
#include <bpf/bpf_helpers.h>
#include <linux/ip.h>
#include <linux/udp.h>
#include <linux/tcp.h>
#include <linux/if_ether.h>
#include <stdbool.h>
#include <stdio.h>
#include <string.h>

#include <linux/in.h>
#include <sys/socket.h>

#define ETH_P_IPV4 0x0800
#define IPPROTO_SRV4 201
#define IPPROTO_IPV4 0X0800
#define IPHDR_LEN 20
#define TLV_TIMEOUT_TYPE 101 
#ifndef memcpy
#define memcpy(dest, src, n) __builtin_memcpy((dest), (src), (n))
#endif


#define TCP_PROTOCOL 6
#define UDP_PROTOCOL 17
#define bpfprint(fmt, ...)                        \
    ({                                             \
        char ____fmt[] = fmt;                      \
        bpf_trace_printk(____fmt, sizeof(____fmt), \
                         ##__VA_ARGS__);           \
    })


#define ipv4_hdrlen(p) ( ( (p)->ihl ) << 2 )
#define ipv4_pllen(p) ( (p)->tot_len - (((p)->ihl)<<2) )
#define srv4_hdrlen(p) ( sizeof(struct ip4_srh_t) + (p)->hdrlen*4 )

		
struct ip4_srh_t 
{
  __u8 nexthdr;
  __u8 hdrlen;
  __u8 type;
  __u8 sl;
  __u8 last_entry;
  __u8 flags;
  __u16 tag;
  __u32 seglist[0];
};

struct bpf_map_def SEC("maps") sidmap = 
{
  .type = BPF_MAP_TYPE_ARRAY,
  .key_size = sizeof(__u32),
  .value_size = sizeof(struct SID),
  .max_entries = MAX_SID_NUM,
};

struct bpf_map_def SEC("maps") pathmap = 
{
  .type = BPF_MAP_TYPE_ARRAY,
  .key_size = sizeof(__u32),
  .value_size = sizeof(struct PATH),
  .max_entries = MAX_PATH_NUM,
};

struct bpf_map_def SEC("maps") flowmap = 
{
  .type = BPF_MAP_TYPE_ARRAY,
  .key_size = sizeof(__u32),
  .value_size = sizeof(struct FLOW),
  .max_entries = MAX_FLOW_NUM,
};

struct bpf_map_def SEC("maps") redirect_map = 
{
  .type = BPF_MAP_TYPE_ARRAY,
  .key_size = sizeof(__u32),
  .value_size = sizeof(struct REDIRECT),
  .max_entries = MAX_REDIRECT_NUM,
};

struct bpf_map_def SEC("maps") macmap = 
{
  .type = BPF_MAP_TYPE_ARRAY,
  .key_size = sizeof(__u32),
  .value_size = sizeof(struct MAC),
  .max_entries = 2,
};

struct bpf_map_def SEC("maps") arpmap = {
	.type = BPF_MAP_TYPE_HASH,
	.key_size = sizeof(__u32),
	.value_size = sizeof(struct ARP),
	.max_entries = MAX_ARP_NUM,
};


static __always_inline __u16 csum_fold_helper(__u32 csum)
{
    __u32 sum;
    sum = (csum >> 16) + (csum & 0xffff);
    sum += (sum >> 16);
    return ~sum;
}

__u16
checksum (__u16 *addr, int len)
{
	int nleft = IPHDR_LEN;
	int sum = 0;
	__u16 *w = addr;
	__u16 answer = 0;

	while (nleft > 1) {
		sum += *w++;
		nleft -= sizeof (__u16);
	}

	if (nleft == 1) {
		*(__u8 *) (&answer) = *(__u8 *) w;
		sum += answer;
	}

	sum = (sum >> 16) + (sum & 0xFFFF);
	sum += (sum >> 16);
	answer = ~sum;
	return (answer);
}

__u16
ip_checksum(struct iphdr *p_iphdr)
{
	p_iphdr->check = 0;
	int iphdrlen = p_iphdr->ihl * 4;

	return checksum ((__u16 *) p_iphdr, iphdrlen);
}

static __always_inline int ip_decrease_ttl(struct iphdr *iph)
{
	__u32 check = (__u32)iph->check;

	check += (__u32)bpf_htons(0x0100);
	iph->check = (__sum16)(check + (check >= 0xFFFF));
	return --iph->ttl;
}

static __always_inline int 
get_net_mac(__u8 mac[ETH_ALEN])
{
	__u32 key = 0;
	struct MAC *p_mac = bpf_map_lookup_elem(&macmap, &key);
	if (p_mac)
	{
		memcpy(mac, p_mac->mac, ETH_ALEN);
		return 0;
	}
	return -1;
}

int 
srv4_ROUTER_process(struct xdp_md *ctx)
{
	void *data_end = (void *)(long)ctx->data_end;
	void *data = (void *)(long)ctx->data;
	struct ethhdr *eth = data;
	struct iphdr *iph;
	struct ARP *pitem = NULL;
	unsigned int daddr = 0;
	__u16 h_proto;
	__u64 nh_off;

	nh_off = sizeof(*eth);
	if (data + nh_off > data_end) {
		return XDP_DROP;
	}

	h_proto = eth->h_proto;
	if (h_proto != bpf_htons(ETH_P_IP)) {
		return XDP_PASS;
	}

	iph = data + nh_off;

	if (iph + 1 > data_end) {
		return XDP_DROP;
	}

	daddr = iph->daddr;

	pitem = bpf_map_lookup_elem(&arpmap, &daddr);
	if (pitem) {
		ip_decrease_ttl(iph);
		memcpy(eth->h_dest, pitem->eth_dest, ETH_ALEN);
		memcpy(eth->h_source, pitem->eth_source, ETH_ALEN);
		return bpf_redirect(pitem->ifindex, 0);
	}
	else 
	{
		bpf_printk("--- ERROR: Failed to find ARP entry! daddr = %x \n", daddr);
		return XDP_PASS;
	}
}


// END指令处理
static int
srv4_END_process(struct xdp_md *ctx)
{
	struct iphdr* p_org_iphdr;
	void *data_end = (void *)(long)ctx->data_end;
	void *data = (void *)(long)ctx->data;
	p_org_iphdr = (void *)(data + sizeof(struct ethhdr));
	if (p_org_iphdr + 1 > data_end)
		goto out;
	if (p_org_iphdr->protocol != IPPROTO_SRV4)
	{
		goto out;
	}
	struct ip4_srh_t* p_srh4;
	if ((void*)p_org_iphdr+ipv4_hdrlen(p_org_iphdr) + sizeof(struct ip4_srh_t) > (void*)(long)ctx->data_end)
	{
		goto out;
	}

	p_srh4 = (struct ip4_srh_t*) ((char*)p_org_iphdr + ipv4_hdrlen(p_org_iphdr));
	if (p_srh4+1 > data_end)
		goto out;
	if (p_srh4->sl < 2)
	{
		goto out;
	}

	char old_iphdr[IPHDR_LEN];
	memcpy(old_iphdr, (const void*)p_org_iphdr, IPHDR_LEN);

	__u32* seg;
	seg = (__u32*) (p_srh4 + 1);
	p_srh4->sl --;
	__u32 addr;
	addr =  *(__u32*)(seg+p_srh4->sl);
	if ((__u32*)(seg+p_srh4->sl)+1>data_end)
		goto out;
	p_org_iphdr->daddr = addr;
	__u16 seed = ((struct iphdr*)old_iphdr)->check;
	__u32 csum = bpf_csum_diff((__be32 *)old_iphdr, IPHDR_LEN, (__be32 *)p_org_iphdr, IPHDR_LEN, ~seed);
	p_org_iphdr->check = csum_fold_helper(csum);

	int action;
	action = srv4_ROUTER_process(ctx);
	return action;

out:
	return XDP_PASS;
}

// DENCAP指令处理
static int
srv4_DENCAP_process(struct xdp_md *ctx)
{
	struct iphdr* p_org_iphdr;
	void *data_end = (void *)(long)ctx->data_end;
	void *data = (void *)(long)ctx->data;
	struct ethhdr *ehdr = data;
	if (ehdr + 1 > data_end) // bounds checking
		goto out;
	p_org_iphdr = (void *)(data + sizeof(struct ethhdr));
	if (p_org_iphdr + 1 > data_end)
		goto out;
	int iphdr_len = ipv4_hdrlen(p_org_iphdr);
	if (iphdr_len != 20)
		goto out;
	iphdr_len = 20;
	struct ip4_srh_t* p_srh4;
	if ((char*)p_org_iphdr+iphdr_len > data_end)
	{
		goto out;
	}

	p_srh4 = (struct ip4_srh_t*) ((char*)p_org_iphdr + iphdr_len);
	if (p_srh4+1 > data_end)
		goto out;
	if (p_srh4->sl != 1)
	{
		goto out;
	}

	int srhdr_len = srv4_hdrlen(p_srh4);

	//1,update dest IP
	if ((__u32*)(p_srh4+1)+4>data_end)
		goto out;
	void *p;
	p = (void*)p_srh4 + sizeof(struct ip4_srh_t);
	p_org_iphdr->daddr = *(__u32*) p;
	p_srh4->sl --;
	p_org_iphdr->daddr = bpf_ntohl(p_org_iphdr->daddr);

	//2,copy ipv4 header
	if ((void*)p_org_iphdr + iphdr_len + srhdr_len > data_end)
		goto out;
	char* p_payload = (char*) ((char*)p_org_iphdr + iphdr_len + srhdr_len);
	char old_iphdr[20];
	char protocol = p_srh4->nexthdr;
	memcpy(old_iphdr, (const void*)p_org_iphdr, iphdr_len);
	if ((long)p_payload-iphdr_len < sizeof(struct iphdr))
		goto out;
	struct iphdr* p_new_iphdr = (struct iphdr*) (p_payload - iphdr_len);
	if (iphdr_len != 20)
		goto out;
	if ((char*)p_org_iphdr+20 > data_end)
		goto out;
	if ((char*)p_new_iphdr+20 > data_end)
		goto out;
	memcpy(p_new_iphdr, (const void *)old_iphdr, 20);
	p_new_iphdr->tot_len -= srhdr_len;
	p_new_iphdr->protocol = protocol;
	//3, copy eth header
	struct ethhdr *p_new_ethhdr;
	p_new_ethhdr = (struct ethhdr*) ((char*)p_new_iphdr - sizeof(struct ethhdr));
	memcpy(p_new_ethhdr, ehdr, sizeof(struct ethhdr));
	__u8 mac[ETH_ALEN];
	if (!get_net_mac(mac))
		memcpy(p_new_ethhdr->h_dest, mac, ETH_ALEN);

	//4,adjust memory
	bpf_xdp_adjust_head(ctx, srhdr_len);
	p_new_iphdr = (struct iphdr*) (ctx->data + sizeof(struct ethhdr));
	if ((void*)p_new_iphdr+IPHDR_LEN > (void*)(long)ctx->data_end)
		goto out;
	memcpy(old_iphdr, (const void*)p_new_iphdr, iphdr_len);
	__u32 checksum = ip_checksum((struct iphdr*)old_iphdr);
	p_new_iphdr->check = checksum;

	int action;
	action = srv4_ROUTER_process(ctx);
	return action;

out:
	return XDP_PASS;
}

// ENCAP指令处理
static int
srv4_ENCAP_process(struct xdp_md *ctx, int pathid, __u16 timeout)
{
	int j;
	struct iphdr* p_org_iphdr;
	void *data_end = (void *)(long)ctx->data_end;
	void *data = (void *)(long)ctx->data;
	struct ethhdr *old_ehdr = data;
	if (old_ehdr + 1 > data_end) // bounds checking
		goto out;
	p_org_iphdr = (struct iphdr*)(data + sizeof(struct ethhdr));
	if (p_org_iphdr + 1 > data_end)
		goto out;
	int iphdr_len = ipv4_hdrlen(p_org_iphdr);
	if (iphdr_len != 20)
		goto out;
	iphdr_len = 20;
	//4.1, 获取path
	bool path_find = false;
	struct PATH* p_path = NULL;
	__u32 key;
	key = (__u32)pathid;
	for (j = 0; j <= MAX_FLOW_NUM; j++) 
	{
		key = (__u32)j;
		p_path = bpf_map_lookup_elem(&pathmap, &key);
		if (!p_path || !p_path->valid)
			continue;
		if (p_path->id == pathid) {
			path_find = true;
			break;
		}
	}
	if (!path_find) {
		bpfprint("---!!! Cannot find path for flow!\n");
		goto out;
	}

	//4.2，扩展报文头空间
	char old_iphdr[20];
	memcpy(old_iphdr, (const void*)p_org_iphdr, iphdr_len);
	char old_ethdr[14];
	memcpy(old_ethdr, (const void*)old_ehdr, 14);

	int extlen = sizeof(struct ip4_srh_t) + p_path->num*4 + 4;
	//bpfprint("--- Extend size is %d, extlen is %d!\n", (long)ctx->data_end-(long)ctx->data, extlen);
	bpf_xdp_adjust_head(ctx, -extlen);
	data_end = (void *)(long)ctx->data_end;

	//4.3, copy eth header
	data = (void *)(long)ctx->data;
	struct ethhdr* p_new_ethhdr = (struct ethhdr*)(long)ctx->data;
	if ((char*)p_new_ethhdr+sizeof(struct ethhdr) > data_end)
		goto out;
	memcpy(p_new_ethhdr, old_ethdr, sizeof(struct ethhdr));
	p_new_ethhdr->h_dest[5] = 0xb5;

	//4.4, copy ipv4 header
	struct iphdr* p_new_iphdr = (struct iphdr*) ((char*)p_new_ethhdr + sizeof(struct ethhdr));
	if ((char*)p_new_iphdr+20 > data_end)
		goto out;
	memcpy((char*)(p_new_iphdr), old_iphdr, 20);
	p_new_iphdr->protocol = IPPROTO_SRV4;
	__u32 daddr = p_path->ipv4[0];
	p_new_iphdr->daddr = bpf_htonl(daddr);
	p_new_iphdr->tot_len += extlen;
	__u16 seed = ((struct iphdr*)old_iphdr)->check;
	__u32 csum = bpf_csum_diff((__be32 *)old_iphdr, iphdr_len, (__be32 *)p_new_iphdr, iphdr_len, ~seed);
	p_new_iphdr->check = csum_fold_helper(csum);

	//4.5，构建SRHv4
	struct ip4_srh_t* p_ip4_srh = (struct ip4_srh_t*) ((char*)(p_new_ethhdr+1)+iphdr_len);
	if ((char*)p_ip4_srh+sizeof(struct ip4_srh_t) > data_end)
		goto out;
	p_ip4_srh->nexthdr = ((struct iphdr*)(old_iphdr))->protocol;
	p_ip4_srh->hdrlen = p_path->num + 1;
	if (timeout > 0)
		p_ip4_srh->hdrlen += sizeof(struct TLV_TIMEOUT)/4;
	p_ip4_srh->type = 4;
	p_ip4_srh->sl = p_path->num;
	p_ip4_srh->last_entry = p_path->num + 1;
	p_ip4_srh->flags = 0;
	p_ip4_srh->tag = 0;

	//4.6, copy segment list
	if (p_path->num > MAX_NODE_NUM)
	{
		bpfprint("The nodes number of path %d is out of %d\n", pathid, MAX_NODE_NUM);
		goto out;
	}
	__u32 *p_seg, addr2;
	p_seg = (__u32 *) ((char*)p_ip4_srh + sizeof(struct ip4_srh_t));
	if ((char*)p_seg+4 > data_end)
		goto out;
	if ((char*)(&(((struct iphdr*)(p_new_ethhdr+1))->daddr)) > data_end)
		goto out;
	daddr = ((struct iphdr*)(old_iphdr))->daddr;
	*p_seg = bpf_htonl( daddr );
	p_seg ++;
	int k;
	for (int j=0; j<MAX_NODE_NUM; j++)
	{
		k = MAX_NODE_NUM-1-j;
		if (p_path->ipv4[k] == 0)
			continue;
		daddr = p_path->ipv4[k];
		addr2 = bpf_htonl(daddr);
		if ((char*)p_seg+4 > data_end)
			goto out;
		*p_seg = addr2;
		p_seg ++;
	}

	int action;
	action = srv4_ROUTER_process(ctx);
	return action;

out:
	return XDP_PASS;
}


SEC("srv4-process")
int xdp_srv4_process(struct xdp_md *ctx) 
{
	//---------------- 预处理 ----------------------
	volatile struct iphdr* p_org_iphdr;
	volatile struct ethhdr old_ehdr;
	void *data_end = (void *)(long)ctx->data_end;
	void *data = (void *)(long)ctx->data;
	int action;

	struct ethhdr *ehdr = data;
	if (ehdr + 1 > data_end) // bounds checking
		goto out;
	old_ehdr = *ehdr;

	if (bpf_ntohs(ehdr->h_proto) != ETH_P_IPV4) {
		goto out;
	}
	p_org_iphdr = (void *)(ehdr + 1);
	if (p_org_iphdr + 1 > data_end)
		goto out;
	if (p_org_iphdr->protocol != IPPROTO_SRV4)
		goto out;

	//---------- 获取目的IP地址，判断是否触发SID ------
	__u8 sid_flag = PASS;
	__u32 key, addr;
		addr = p_org_iphdr->daddr;
		addr = bpf_htonl(addr);
	int j;
	for (j = 0; j <= MAX_SID_NUM; j++) 
	{
		__u32 key = (__u32)j;
		struct SID *psid = bpf_map_lookup_elem(&sidmap, &key);
		if (!psid)
			continue;
		if (!psid->valid)
			continue;
		if (bpf_htonl(psid->ipv4) == p_org_iphdr->daddr)
		{
			sid_flag = psid->action;
			break;
		}
	}
	//---------- __END__ ---------------------------

	//---------- End指令处理 -----------------------
	if (sid_flag == END && p_org_iphdr->protocol == IPPROTO_SRV4)
	{
		action = srv4_END_process(ctx);
		return action;
	}
	//--------- __END__ ----------------------------  


	//---------- Decap处理 ------------------------
	if (sid_flag == DENCAP && p_org_iphdr->protocol == IPPROTO_SRV4)
	{
		action = srv4_DENCAP_process(ctx);
		return action;
	}
	//--------- __END__ ----------------------------  

	//---------- 获取五元组，匹配流 ------------------
	__u32 src_ip, dst_ip;
	__u16 src_port, dst_port;
	__u8 proto;
	struct udphdr *p_udphdr;

	//1，判断是否为UDP或TCP报文
	if (p_org_iphdr->protocol!=UDP_PROTOCOL && p_org_iphdr->protocol!=TCP_PROTOCOL) {
		goto out;
	}

	//2，获取五元组信息
	src_ip = bpf_ntohl(p_org_iphdr->saddr);
	dst_ip = bpf_ntohl(p_org_iphdr->daddr);
	int iphdr_len = ipv4_hdrlen(p_org_iphdr);
		if (iphdr_len != 20)
			goto out;
		iphdr_len = 20;
	p_udphdr = (struct udphdr*)((char*)p_org_iphdr + iphdr_len);
		if ((char*)p_udphdr+sizeof(struct udphdr) > data_end)
			goto out;
	src_port = bpf_ntohs(p_udphdr->source);
	dst_port = bpf_ntohs(p_udphdr->dest);
	proto = p_org_iphdr->protocol;

	//3，查询匹配FLOW
	int pathid = -1;
	__u16 timeout = 0;
	for (j = 0; j <= MAX_FLOW_NUM; j++) 
	{
		key = (__u32)j;
		struct FLOW* pflow = bpf_map_lookup_elem(&flowmap, &key);
		if (!pflow || !pflow->valid)
			continue;
		if ( (pflow->src_ipv4==src_ip) &&
				(pflow->dst_ipv4==dst_ip) &&
				(pflow->src_port==0 || pflow->src_port==src_port) &&
				(pflow->dst_port==0 || pflow->dst_port==dst_port) )
		{
			pathid = pflow->id;
			timeout = (__u16)(pflow->timeout);
			break;
		}
	}
	//4，匹配命中，封装SRHv4头
	if (pathid>=0) 
	{
		action = srv4_ENCAP_process(ctx, pathid, timeout);
		return action;
	}
	//---------- Encaps处理 -----------------------

out:
	return XDP_PASS;
}

char _license[] SEC("license") = "GPL";





